#pragma once
#include "decode.cuh"
using namespace cv;

namespace SLCuda
{
    __device__ int getPhaseDevice(double sinPhase, double cosPhase)
    {
        if (sinPhase > 0 && cosPhase > 0)
            return 0;
        if (sinPhase > 0 && cosPhase < 0)
            return 1;
        if (sinPhase < 0 && cosPhase < 0)
            return 2;
        if (sinPhase < 0 && cosPhase > 0)
            return 3;

        return 0;
    }

    __global__ void assginPointer(uchar **imagesCuda, uchar *image, int n)
    {
        imagesCuda[n] = image;
    }

    __global__ void decodePhaseShiftKernel(uchar **imagesCuda, float *decodeImage, float *uncertainImage, int imagesNo, int length, float phaseShiftCutoff)
    {
        double sinSum = 0;
        double cosSum = 0;
        // int tid = blockIdx.z * (gridDim.x * gridDim.y) * (blockDim.x * blockDim.y * blockDim.z) + blockIdx.y * gridDim.x * (blockDim.x * blockDim.y * blockDim.z) + blockIdx.x * (blockDim.x * blockDim.y * blockDim.z) + threadIdx.z * (blockDim.x * blockDim.y) + threadIdx.y * blockDim.x + threadIdx.x;
        int tid = blockIdx.x * blockDim.x * blockDim.y + threadIdx.y * blockDim.x + threadIdx.x;
        // if (tid < length)
        // {
        for (int i = 0; i < imagesNo; i++)
        {
            double x = 2 * PI * double(i) / imagesNo;
            sinSum += (double)(imagesCuda[i])[tid] * -sin(x);
            cosSum += (double)(imagesCuda[i])[tid] * cos(x);
        }

        double tempPhase = (double)(atan(sinSum / cosSum));
        int phase = getPhaseDevice(sinSum, cosSum);
        double add;
        if (phase == 0)
            add = 0;
        if (phase == 1)
            add = PI;
        if (phase == 2)
            add = PI;
        if (phase == 3)
            add = 2 * PI;
        // decodeImage[tid] = (float)(100.0f);
        decodeImage[tid] = (float)(add + tempPhase);
        // decodeImage[tid] = (imagesCuda[0][tid] + imagesCuda[1][tid]) / 2;
        float b = (float)(sqrt(sinSum * sinSum + cosSum * cosSum));
        uncertainImage[tid] = b < phaseShiftCutoff ? 1.0f : 0;
        // BitCodeList bitCodeList = generateGrayCodeDevice(5);
        // }
    };

    Mat decodePhaseShiftCuda(SL::Images &images, Size size, int n, float phaseShiftCutoff)
    {
        int length = size.height * size.width;
        int memSizeFloat = size.height * size.width * sizeof(float);
        int memSizeUchar = size.height * size.width * sizeof(uchar);
        int imagesNo = images.size();

        uchar **imagesCuda;
        cudaMalloc((void **)&imagesCuda, imagesNo * length * sizeof(uchar *));

        for (int i = 0; i < imagesNo; ++i)
        {
            uchar *imageDevice;
            cudaMalloc((void **)&imageDevice, memSizeUchar);
            cudaMemcpy(imageDevice, images[i].data, memSizeUchar, cudaMemcpyHostToDevice);
            assginPointer<<<1, 1>>>(imagesCuda, imageDevice, i);
        }

        float *decodeImageDevice, *uncertainImageDevice;
        cudaMalloc((void **)&decodeImageDevice, memSizeFloat);
        cudaMalloc((void **)&uncertainImageDevice, memSizeFloat);

        dim3 block(1 + (length / (32 * 32 + 1)), 1, 1); // grid
        dim3 thread(32, 32, 1);
        decodePhaseShiftKernel<<<block, thread>>>(imagesCuda, decodeImageDevice, uncertainImageDevice, imagesNo, length, phaseShiftCutoff);
        cudaDeviceSynchronize();

        Mat decodeImageHost = Mat::zeros(size.height, size.width, CV_32FC1);
        Mat uncertainImageHost = Mat::zeros(size.height, size.width, CV_32FC1);
        cudaMemcpy(decodeImageHost.data, decodeImageDevice, length * sizeof(float), cudaMemcpyDeviceToHost);
        cudaMemcpy(uncertainImageHost.data, uncertainImageDevice, length * sizeof(float), cudaMemcpyDeviceToHost);
        std::vector<Mat> channels = {decodeImageHost, uncertainImageHost};
        Mat decodeImage = Mat::zeros(size.height, size.width, CV_32FC2);
        merge(channels, decodeImage);
        return decodeImage;
    }

    __global__ void decodeGrayCodeKernel(uchar **imagesCuda, ushort *decodeImage, ushort *uncertainImage, int n, float cutoffPercentage)
    {
        int min = 256;
        int max = 0;
        int tid = blockIdx.x * blockDim.x * blockDim.y + threadIdx.y * blockDim.x + threadIdx.x;
        for (int i = 0; i < n + 2; i++)
        {
            if (imagesCuda[i][tid] < min)
                min = (int)imagesCuda[i][tid];
            if (imagesCuda[i][tid] > max)
                max = (int)imagesCuda[i][tid];
        }
        if (max - min < 10)
        {
            uncertainImage[tid] = 1;
            return;
        }
        for (int i = 2; i < n + 2; i++)
        {
            if (imagesCuda[i][tid] > cutoffPercentage * (max - min) + min)
            {
                decodeImage[tid] += 1 << (n - (i - 2) - 1);
            }
        }
    }

    Mat decodeGrayCodeCuda(SL::Images &images, Size size, int n, float cutoffPercentage)
    {
        int length = size.height * size.width;
        int memSizeUshort = size.height * size.width * sizeof(ushort);
        int memSizeUchar = size.height * size.width * sizeof(uchar);
        int imagesNo = images.size();

        uchar **imagesCuda;
        cudaMalloc((void **)&imagesCuda, imagesNo * length * sizeof(uchar *));
        for (int i = 0; i < imagesNo; ++i)
        {
            uchar *imageDevice;
            cudaMalloc((void **)&imageDevice, memSizeUchar);
            cudaMemcpy(imageDevice, images[i].data, memSizeUchar, cudaMemcpyHostToDevice);
            assginPointer<<<1, 1>>>(imagesCuda, imageDevice, i);
        }

        ushort *decodeImageDevice, *uncertainImageDevice;
        cudaMalloc((void **)&decodeImageDevice, memSizeUshort);
        cudaMalloc((void **)&uncertainImageDevice, memSizeUshort);

        dim3 block(1 + (length / (32 * 32 + 1)), 1, 1); // grid
        dim3 thread(32, 32, 1);
        decodeGrayCodeKernel<<<block, thread>>>(imagesCuda, decodeImageDevice, uncertainImageDevice, n, cutoffPercentage);
        cudaDeviceSynchronize();

        Mat decodeImageHost = Mat::zeros(size.height, size.width, CV_16UC1);
        Mat uncertainImageHost = Mat::zeros(size.height, size.width, CV_16UC1);
        cudaMemcpy(decodeImageHost.data, decodeImageDevice, length * sizeof(ushort), cudaMemcpyDeviceToHost);
        cudaMemcpy(uncertainImageHost.data, uncertainImageDevice, length * sizeof(ushort), cudaMemcpyDeviceToHost);

        Mat decodeImage = Mat::zeros(size.height, size.width, CV_16UC2);
        std::vector<Mat> channels = {decodeImageHost, uncertainImageHost};
        merge(channels, decodeImage);
        return decodeImage;
    }

    BitCodeList generateGrayCode(int n)
    {
        BitCodeList bitCodeList;
        bitCodeList.resize(1i64 << n);
        if (n == 1)
        {
            BitCode bitCode0 = {0};
            BitCode bitCode1 = {1};
            bitCodeList[0] = bitCode0;
            bitCodeList[1] = bitCode1;
            return bitCodeList;
        }

        BitCodeList preBitCodeList = generateGrayCode(n - 1);
        int preBitCodeListNo = 1 << (n - 1);
        for (int i = 0; i < preBitCodeListNo; i++)
        {
            BitCode bitcode0 = {0};
            bitcode0.insert(bitcode0.end(), preBitCodeList[i].begin(), preBitCodeList[i].end());
            bitCodeList[i] = bitcode0;
            BitCode bitcode1 = {1};
            bitcode1.insert(bitcode1.end(), preBitCodeList[preBitCodeListNo - i - 1].begin(), preBitCodeList[preBitCodeListNo - i - 1].end());
            bitCodeList[preBitCodeListNo + i] = bitcode1;
        }
        return bitCodeList;
    }

    std::map<int, int> generateBitCodeMap(int grayCodeN)
    {
        BitCodeList bitCodeList = generateGrayCode(grayCodeN);
        std::map<int, int> bitCodeMap;
        for (int i = 0; i < (1 << grayCodeN); i++)
        {
            int x = 0;
            for (int n = 0; n < grayCodeN; n++)
            {
                x += bitCodeList[i][n] != 0 ? (1 << (grayCodeN - n - 1)) : 0;
            }
            bitCodeMap.insert(std::pair<int, int>(x, i));
        }
        return bitCodeMap;
    };

    __global__ void decodeGrayCodeWithPhaseShiftKernel(ushort *grayDecodeImageDevice, ushort *grayUncertainImageDevice, ushort *unwrapGrayDecodeImageDevice, ushort *unwrapGrayUncertainImageDevice, float *phaseShiftDecodeImageDevice, float *phaseShiftUncertainImageDevice, int *bitCodeMapFirstDevice, int *bitCodeMapSecondDevice, int *unwrapBitCodeMapFirstDevice, int *unwrapBitCodeMapSecondDevice, float *decodeImage, int grayCodeN, int stripSize)
    {
        int tid = blockIdx.x * blockDim.x * blockDim.y + threadIdx.y * blockDim.x + threadIdx.x;
        ushort bitCode = grayDecodeImageDevice[tid];
        float uncertain = grayUncertainImageDevice[tid] + phaseShiftUncertainImageDevice[tid];
        ushort unwrapBitCode = unwrapGrayDecodeImageDevice[tid];

        int bitCodeSecond = -1;
        for (int i = 0; i < (1 << grayCodeN); i++)
        {
            if (bitCode == bitCodeMapFirstDevice[i])
            {
                bitCodeSecond = bitCodeMapSecondDevice[i];
                break;
            }
        }

        int unwrapBitCodeSecond = -1;
        for (int i = 0; i < (1 << (grayCodeN + 1)); i++)
        {
            if (unwrapBitCode == unwrapBitCodeMapFirstDevice[i])
            {
                unwrapBitCodeSecond = unwrapBitCodeMapSecondDevice[i];
                break;
            }
        }

        if (uncertain == 0 && bitCodeSecond != -1 && unwrapBitCodeSecond != -1)
        {
            float phase = phaseShiftDecodeImageDevice[tid];
            float finalPos = phase / (2 * PI) * stripSize;
            if (phase < 0.5 * PI)
            {
                float newPhase = phase + PI;
                int grayId = unwrapBitCodeSecond;
                int startStripId = grayId % 2 == 0 ? grayId - 1 : grayId;
                decodeImage[tid] = newPhase / (2 * PI) * stripSize + (startStripId)*stripSize / 2.0f;
            }
            else if (phase > 1.5 * PI)
            {
                float newPhase = phase - PI;
                int grayId = unwrapBitCodeSecond;
                int startStripId = grayId % 2 == 0 ? grayId - 1 : grayId;
                decodeImage[tid] = newPhase / (2 * PI) * stripSize + startStripId * stripSize / 2.0f;
            }
            else
            {
                decodeImage[tid] = finalPos + bitCodeSecond * stripSize;
            }
        }
        else
        {
            decodeImage[tid] = NAN;
        }
    }

    Mat decodeGrayCodeWithPhaseShiftCuda(Mat grayDecodeImage, Mat unwrapGrayDecodeImage, Mat phaseShiftDecodeImage, Size imageSize, Size projectorSize, int grayCodeN)
    {
        int length = imageSize.height * imageSize.width;
        int memSizeUshort = imageSize.height * imageSize.width * sizeof(ushort);
        int memSizeUchar = imageSize.height * imageSize.width * sizeof(uchar);
        int memSizeFloat = imageSize.height * imageSize.width * sizeof(float);
        int stripSize = projectorSize.width / (1 << grayCodeN);

        ushort *grayDecodeImageDevice, *grayUncertainImageDevice;
        Mat grayCodeChannels[2];
        split(grayDecodeImage, grayCodeChannels);
        cudaMalloc((void **)&grayDecodeImageDevice, memSizeUshort);
        cudaMemcpy(grayDecodeImageDevice, grayCodeChannels[0].data, memSizeUshort, cudaMemcpyHostToDevice);
        cudaMalloc((void **)&grayUncertainImageDevice, memSizeUshort);
        cudaMemcpy(grayUncertainImageDevice, grayCodeChannels[1].data, memSizeUshort, cudaMemcpyHostToDevice);

        ushort *unwrapGrayDecodeImageDevice, *unwrapGrayUncertainImageDevice;
        Mat unwrapGrayCodeChannels[2];
        split(unwrapGrayDecodeImage, unwrapGrayCodeChannels);
        cudaMalloc((void **)&unwrapGrayDecodeImageDevice, memSizeUshort);
        cudaMemcpy(unwrapGrayDecodeImageDevice, unwrapGrayCodeChannels[0].data, memSizeUshort, cudaMemcpyHostToDevice);
        cudaMalloc((void **)&unwrapGrayUncertainImageDevice, memSizeUshort);
        cudaMemcpy(unwrapGrayUncertainImageDevice, unwrapGrayCodeChannels[1].data, memSizeUshort, cudaMemcpyHostToDevice);

        float *phaseShiftDecodeImageDevice, *phaseShiftUncertainImageDevice;
        Mat phaseShiftCodeChannels[2];
        split(phaseShiftDecodeImage, phaseShiftCodeChannels);
        cudaMalloc((void **)&phaseShiftDecodeImageDevice, memSizeFloat);
        cudaMemcpy(phaseShiftDecodeImageDevice, phaseShiftCodeChannels[0].data, memSizeFloat, cudaMemcpyHostToDevice);
        cudaMalloc((void **)&phaseShiftUncertainImageDevice, memSizeFloat);
        cudaMemcpy(phaseShiftUncertainImageDevice, phaseShiftCodeChannels[1].data, memSizeFloat, cudaMemcpyHostToDevice);

        std::map<int, int> bitCodeMap = generateBitCodeMap(grayCodeN);
        std::vector<int> bitCodeMapFirst, bitCodeMapSecond;
        for (auto it = bitCodeMap.begin(); it != bitCodeMap.end(); ++it)
        {
            bitCodeMapFirst.push_back(it->first);
            bitCodeMapSecond.push_back(it->second);
        }
        int *bitCodeMapFirstDevice, *bitCodeMapSecondDevice;
        cudaMalloc((void **)&bitCodeMapFirstDevice, bitCodeMap.size() * sizeof(int));
        cudaMalloc((void **)&bitCodeMapSecondDevice, bitCodeMap.size() * sizeof(int));
        cudaMemcpy(bitCodeMapFirstDevice, bitCodeMapFirst.data(), bitCodeMap.size() * sizeof(int), cudaMemcpyHostToDevice);
        cudaMemcpy(bitCodeMapSecondDevice, bitCodeMapSecond.data(), bitCodeMap.size() * sizeof(int), cudaMemcpyHostToDevice);

        std::map<int, int> unwrapBitCodeMap = generateBitCodeMap(grayCodeN + 1);
        std::vector<int> unwrapBitCodeMapFirst, unwrapBitCodeMapSecond;
        for (auto it = unwrapBitCodeMap.begin(); it != unwrapBitCodeMap.end(); ++it)
        {
            unwrapBitCodeMapFirst.push_back(it->first);
            unwrapBitCodeMapSecond.push_back(it->second);
        }
        int *unwrapBitCodeMapFirstDevice, *unwrapBitCodeMapSecondDevice;
        cudaMalloc((void **)&unwrapBitCodeMapFirstDevice, unwrapBitCodeMap.size() * sizeof(int));
        cudaMalloc((void **)&unwrapBitCodeMapSecondDevice, unwrapBitCodeMap.size() * sizeof(int));
        cudaMemcpy(unwrapBitCodeMapFirstDevice, unwrapBitCodeMapFirst.data(), unwrapBitCodeMap.size() * sizeof(int), cudaMemcpyHostToDevice);
        cudaMemcpy(unwrapBitCodeMapSecondDevice, unwrapBitCodeMapSecond.data(), unwrapBitCodeMap.size() * sizeof(int), cudaMemcpyHostToDevice);

        float *decodeImageDevice;
        cudaMalloc((void **)&decodeImageDevice, memSizeFloat);

        dim3 block(1 + (length / (16 * 16 + 1)), 1, 1);
        dim3 thread(16, 16, 1);
        decodeGrayCodeWithPhaseShiftKernel<<<block, thread>>>(grayDecodeImageDevice, grayUncertainImageDevice, unwrapGrayDecodeImageDevice, unwrapGrayUncertainImageDevice, phaseShiftDecodeImageDevice, phaseShiftUncertainImageDevice, bitCodeMapFirstDevice, bitCodeMapSecondDevice, unwrapBitCodeMapFirstDevice, unwrapBitCodeMapSecondDevice, decodeImageDevice, grayCodeN, stripSize);

        Mat decodeImage = Mat::zeros(imageSize.height, imageSize.width, CV_32FC1);
        cudaMemcpy(decodeImage.data, decodeImageDevice, memSizeFloat, cudaMemcpyDeviceToHost);
        // imshow("", decodeImage);
        // waitKey(0);
        return decodeImage;
    }

}
